{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "49a37d56-e6f4-4c23-99e6-577ba604cf31",
   "metadata": {},
   "source": [
    "# L3: Preparing for on-device deployment\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "47a6b277-2dc7-44d4-9571-36635d3aae68",
   "metadata": {},
   "source": [
    "<p style=\"background-color:#fff6e4; padding:15px; border-width:3px; border-color:#f5ecda; border-style:solid; border-radius:6px\"> ⏳ <b>Note <code>(Kernel Starting)</code>:</b> This notebook takes about 30 seconds to be ready to use. You may start and watch the video while you wait.</p>\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "1a920ffb-05b1-4550-905e-df45832621c2",
   "metadata": {},
   "source": [
    "## Capture trained model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "51717f2c-bd2e-4d3e-b514-c831ac646120",
   "metadata": {},
   "outputs": [],
   "source": [
    "from qai_hub_models.models.ffnet_40s import Model as FFNet_40s\n",
    "\n",
    "# Load from pre-trained weights\n",
    "ffnet_40s = FFNet_40s.from_pretrained()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "188e2b3b-1f3b-423d-950d-6d095bf42601",
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch\n",
    "input_shape = (1, 3, 1024, 2048)\n",
    "example_inputs = torch.rand(input_shape)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8d245965-93f9-4ca9-b322-b5d909c798b6",
   "metadata": {},
   "outputs": [],
   "source": [
    "traced_model = torch.jit.trace(ffnet_40s, example_inputs)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d4f6c00f-a7b0-4061-b933-46b1ea05405c",
   "metadata": {},
   "outputs": [],
   "source": [
    "traced_model"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "8dfd3dd5-1b7e-4438-92dd-e673c3d9fcd2",
   "metadata": {},
   "source": [
    "## Compile for device"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d8d445a7-223b-4721-986b-07abba7cdb7b",
   "metadata": {},
   "source": [
    "<p style=\"background-color:#fff6ff; padding:15px; border-width:3px; border-color:#efe6ef; border-style:solid; border-radius:6px\"> 💻 &nbsp; <b>Access Utils File and Helper Functions:</b> To access the files for this notebook, 1) click on the <em>\"File\"</em> option on the top menu of the notebook and then 2) click on <em>\"Open\"</em>. For more help, please see the <em>\"Appendix - Tips and Help\"</em> Lesson.</p>"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "2a69895b-ac9c-4e80-b2f7-f6205cc7a694",
   "metadata": {},
   "outputs": [],
   "source": [
    "import qai_hub\n",
    "import qai_hub_models\n",
    "\n",
    "from utils import get_ai_hub_api_token\n",
    "ai_hub_api_token = get_ai_hub_api_token()\n",
    "\n",
    "!qai-hub configure --api_token $ai_hub_api_token"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e6b1f780-b534-488d-97cb-11b362944e35",
   "metadata": {},
   "outputs": [],
   "source": [
    "for device in qai_hub.get_devices():\n",
    "    print(device.name)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "455936c0-67f8-4f26-925c-de58bfe16363",
   "metadata": {},
   "source": [
    "<p style=\"background-color:#fff6e4; padding:15px; border-width:3px; border-color:#f5ecda; border-style:solid; border-radius:6px\"> ⏳ <b>Note:</b> To spread the load across various devices, we are selecting a random device. Feel free to change it to any other device you prefer.</p>"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "2c9f33e5-4039-4b2c-8eaa-974fc8c4fa4a",
   "metadata": {},
   "outputs": [],
   "source": [
    "devices = [\n",
    "    \"Samsung Galaxy S22 Ultra 5G\",\n",
    "    \"Samsung Galaxy S22 5G\",\n",
    "    \"Samsung Galaxy S22+ 5G\",\n",
    "    \"Samsung Galaxy Tab S8\",\n",
    "    \"Xiaomi 12\",\n",
    "    \"Xiaomi 12 Pro\",\n",
    "    \"Samsung Galaxy S22 5G\",\n",
    "    \"Samsung Galaxy S23\",\n",
    "    \"Samsung Galaxy S23+\",\n",
    "    \"Samsung Galaxy S23 Ultra\",\n",
    "    \"Samsung Galaxy S24\",\n",
    "    \"Samsung Galaxy S24 Ultra\",\n",
    "    \"Samsung Galaxy S24+\",\n",
    "]\n",
    "\n",
    "import random\n",
    "selected_device = random.choice(devices)\n",
    "print(selected_device)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "3c77c567-383b-4032-8b34-6331fa3a7c13",
   "metadata": {},
   "outputs": [],
   "source": [
    "device = qai_hub.Device(selected_device)\n",
    "\n",
    "# Compile for target device\n",
    "compile_job = qai_hub.submit_compile_job(\n",
    "    model=traced_model,                        # Traced PyTorch model\n",
    "    input_specs={\"image\": input_shape},        # Input specification\n",
    "    device=device,                             # Device\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "25948044-7b8d-48f9-be1b-8c508942cd21",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Download and save the target model for use on-device\n",
    "target_model = compile_job.get_target_model()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "7d69562e-558c-4be7-99f7-56165ca9ded9",
   "metadata": {},
   "source": [
    "## Exercise: Try different runtimes "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d6a72528-4329-4ec0-950d-f98c7380bac9",
   "metadata": {},
   "outputs": [],
   "source": [
    "compile_options=\"--target_runtime tflite\"                  # Uses TensorFlow Lite\n",
    "compile_options=\"--target_runtime onnx\"                    # Uses ONNX runtime\n",
    "compile_options=\"--target_runtime qnn_lib_aarch64_android\" # Runs with Qualcomm AI Engine\n",
    "\n",
    "compile_job_expt = qai_hub.submit_compile_job(\n",
    "    model=traced_model,                        # Traced PyTorch model\n",
    "    input_specs={\"image\": input_shape},        # Input specification\n",
    "    device=device,                             # Device\n",
    "    options=compile_options,\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "1f855e6e-7146-4b3f-bf11-8711c4b078ed",
   "metadata": {},
   "source": [
    "Expore more compiler options <a href=https://app.aihub.qualcomm.com/docs/hub/compile_examples.html#compiling-pytorch-to-tflite> here</a>."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "16217faa-ce79-47e3-ba78-0c3c1b280b45",
   "metadata": {},
   "source": [
    "## On-Device Performance Profiling"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "204f5f67-1ab9-463b-b826-48324b9780a7",
   "metadata": {},
   "outputs": [],
   "source": [
    "from qai_hub_models.utils.printing import print_profile_metrics_from_job\n",
    "\n",
    "# Choose device\n",
    "device = qai_hub.Device(selected_device)\n",
    "\n",
    "# Runs a performance profile on-device\n",
    "profile_job = qai_hub.submit_profile_job(\n",
    "    model=target_model,                       # Compiled model\n",
    "    device=device,                            # Device\n",
    ")\n",
    "\n",
    "# Print summary\n",
    "profile_data = profile_job.download_profile()\n",
    "print_profile_metrics_from_job(profile_job, profile_data)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "74910fb1-0c98-4f6e-97fa-67cb247eac77",
   "metadata": {},
   "source": [
    "## Exercise: Try different compute units"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "edac1287-2a85-42e6-807e-77662be565d8",
   "metadata": {},
   "outputs": [],
   "source": [
    "profile_options=\"--compute_unit cpu\"     # Use cpu \n",
    "profile_options=\"--compute_unit gpu\"     # Use gpu (with cpu fallback) \n",
    "profile_options=\"--compute_unit npu\"     # Use npu (with cpu fallback) \n",
    "\n",
    "# Runs a performance profile on-device\n",
    "profile_job_expt = qai_hub.submit_profile_job(\n",
    "    model=target_model,                     # Compiled model\n",
    "    device=device,                          # Device\n",
    "    options=profile_options,\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "59a9b6fa-4666-411f-a08c-e7fcecf8be8f",
   "metadata": {},
   "source": [
    "## On-Device Inference"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f427aee3-b141-464a-98b3-d16961575d14",
   "metadata": {},
   "outputs": [],
   "source": [
    "sample_inputs = ffnet_40s.sample_inputs()\n",
    "sample_inputs"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e8b4d2e2-2715-41bf-8e44-de3937e7a7f6",
   "metadata": {},
   "outputs": [],
   "source": [
    "torch_inputs = torch.Tensor(sample_inputs['image'][0])\n",
    "torch_outputs = ffnet_40s(torch_inputs)\n",
    "torch_outputs"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "19ad375b-2482-425e-a756-070fcec905df",
   "metadata": {},
   "outputs": [],
   "source": [
    "inference_job = qai_hub.submit_inference_job(\n",
    "        model=target_model,          # Compiled model\n",
    "        inputs=sample_inputs,        # Sample input\n",
    "        device=device,               # Device\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "49ad79df-bf53-49b9-b785-3348cce74d46",
   "metadata": {},
   "outputs": [],
   "source": [
    "ondevice_outputs = inference_job.download_output_data()\n",
    "ondevice_outputs['output_0']"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "97670114-573c-41a9-8b62-dc9caa98e642",
   "metadata": {},
   "outputs": [],
   "source": [
    "from qai_hub_models.utils.printing import print_inference_metrics\n",
    "print_inference_metrics(inference_job, ondevice_outputs, torch_outputs)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "9fffc80a-3562-4314-934d-832eff2d9a10",
   "metadata": {},
   "source": [
    "## Get ready for deployment!"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "5e1ced93-f86c-4097-bc79-439926e1c6df",
   "metadata": {},
   "outputs": [],
   "source": [
    "target_model = compile_job.get_target_model()\n",
    "_ = target_model.download(\"FFNet_40s.tflite\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7402ad4f-c99a-4980-b462-a93524c55c85",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "194a5621-e8f6-47dc-b103-09c20fb428e6",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c18a91a8-61fc-4690-826b-72660b7b584b",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7aa6e343-a588-4c60-bdf6-3a7857430442",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1cd98f15-1ced-4ab9-bc06-6e16dbc1389c",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.8.19"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
